"""
15_pathfinding_algorithms.py
============================
Pathfinding Algorithms with Neo4j GDS

This module covers:
1. Shortest Path algorithms
2. All Shortest Paths
3. Dijkstra's algorithm
4. A* algorithm
5. Yen's K-Shortest Paths
6. Breadth-First Search

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


@dataclass
class PathResult:
    """Result from pathfinding algorithm."""
    source: str
    target: str
    path: List[str]
    cost: float
    hops: int


class PathfindingAlgorithms:
    """
    Implements pathfinding algorithms using Neo4j GDS.
    
    Pathfinding helps answer questions like:
    - What's the shortest treatment path for a patient?
    - How are two diseases connected through medications?
    - What's the minimum cost treatment plan?
    """
    
    def __init__(self, driver: Driver):
        self.driver = driver
        self.graph_name = "pathfinding_graph"
    
    def create_graph_projection(
        self,
        include_weights: bool = True
    ) -> Dict[str, Any]:
        """Create graph projection for pathfinding."""
        drop_query = f"CALL gds.graph.drop('{self.graph_name}', false)"
        
        if include_weights:
            create_query = """
            CALL gds.graph.project(
                $graph_name,
                ['Patient', 'Disease', 'Medication', 'Doctor'],
                {
                    DIAGNOSED_WITH: {
                        orientation: 'UNDIRECTED',
                        properties: ['severity']
                    },
                    TAKES: {
                        orientation: 'UNDIRECTED',
                        properties: ['dosage']
                    },
                    TREATS: {
                        orientation: 'UNDIRECTED',
                        properties: ['efficacy']
                    },
                    PRESCRIBED_BY: {
                        orientation: 'UNDIRECTED'
                    }
                }
            )
            YIELD graphName, nodeCount, relationshipCount
            RETURN graphName, nodeCount, relationshipCount
            """
        else:
            create_query = """
            CALL gds.graph.project(
                $graph_name,
                ['Patient', 'Disease', 'Medication', 'Doctor'],
                {
                    DIAGNOSED_WITH: {orientation: 'UNDIRECTED'},
                    TAKES: {orientation: 'UNDIRECTED'},
                    TREATS: {orientation: 'UNDIRECTED'},
                    PRESCRIBED_BY: {orientation: 'UNDIRECTED'}
                }
            )
            YIELD graphName, nodeCount, relationshipCount
            RETURN graphName, nodeCount, relationshipCount
            """
        
        with self.driver.session() as session:
            try:
                session.run(drop_query)
            except:
                pass
            
            result = session.run(create_query, graph_name=self.graph_name)
            record = result.single()
            return record.data() if record else None
    
    # =========================================================
    # DIJKSTRA'S SHORTEST PATH
    # =========================================================
    
    def dijkstra_shortest_path(
        self,
        source_id: str,
        target_id: str,
        weight_property: str = None
    ) -> Dict[str, Any]:
        """
        Find shortest path using Dijkstra's algorithm.
        
        Dijkstra finds the path with minimum total weight.
        Without weights, finds path with fewest hops.
        
        Args:
            source_id: Starting node ID
            target_id: Target node ID
            weight_property: Optional edge weight property
            
        Returns:
            Path information including nodes and total cost
            
        Use Cases:
            - Find shortest treatment path
            - Minimum cost medication route
            - Fastest referral chain
        """
        config = {}
        if weight_property:
            config['relationshipWeightProperty'] = weight_property
        
        query = """
        MATCH (source {id: $source_id}), (target {id: $target_id})
        CALL gds.shortestPath.dijkstra.stream($graph_name, {
            sourceNode: source,
            targetNode: target,
            relationshipWeightProperty: $weight_prop
        })
        YIELD index, sourceNode, targetNode, totalCost, nodeIds, costs, path
        WITH totalCost, 
             [nodeId IN nodeIds | gds.util.asNode(nodeId).id] AS nodeIdList,
             [nodeId IN nodeIds | gds.util.asNode(nodeId).name] AS nodeNames,
             costs
        RETURN 
            nodeIdList,
            nodeNames,
            totalCost,
            size(nodeIdList) - 1 AS hops,
            costs
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                source_id=source_id,
                target_id=target_id,
                weight_prop=weight_property
            )
            record = result.single()
            return record.data() if record else None
    
    # =========================================================
    # ALL SHORTEST PATHS
    # =========================================================
    
    def all_shortest_paths(
        self,
        source_id: str,
        target_id: str,
        max_depth: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Find all shortest paths between two nodes.
        
        Returns multiple paths if they have the same length.
        
        Args:
            source_id: Starting node ID
            target_id: Target node ID
            max_depth: Maximum path length
            
        Returns:
            List of all shortest paths
        """
        query = f"""
        MATCH (source {{id: $source_id}}), (target {{id: $target_id}})
        MATCH paths = allShortestPaths((source)-[*..{max_depth}]-(target))
        WITH paths
        RETURN 
            [node IN nodes(paths) | node.id] AS nodeIds,
            [node IN nodes(paths) | node.name] AS nodeNames,
            [rel IN relationships(paths) | type(rel)] AS relationshipTypes,
            length(paths) AS pathLength
        LIMIT 10
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                source_id=source_id,
                target_id=target_id
            )
            return [record.data() for record in result]
    
    # =========================================================
    # YEN'S K-SHORTEST PATHS
    # =========================================================
    
    def k_shortest_paths(
        self,
        source_id: str,
        target_id: str,
        k: int = 3,
        weight_property: str = None
    ) -> List[Dict[str, Any]]:
        """
        Find K shortest paths using Yen's algorithm.
        
        Returns the K best paths, not just the single best.
        Useful for finding alternative routes.
        
        Args:
            source_id: Starting node ID
            target_id: Target node ID
            k: Number of paths to find
            weight_property: Optional weight property
            
        Returns:
            K shortest paths with costs
            
        Use Cases:
            - Alternative treatment options
            - Backup referral paths
            - Multiple medication routes
        """
        query = """
        MATCH (source {id: $source_id}), (target {id: $target_id})
        CALL gds.shortestPath.yens.stream($graph_name, {
            sourceNode: source,
            targetNode: target,
            k: $k,
            relationshipWeightProperty: $weight_prop
        })
        YIELD index, sourceNode, targetNode, totalCost, nodeIds, costs
        WITH index, totalCost,
             [nodeId IN nodeIds | gds.util.asNode(nodeId).id] AS nodeIdList,
             [nodeId IN nodeIds | gds.util.asNode(nodeId).name] AS nodeNames
        RETURN 
            index AS pathIndex,
            nodeIdList,
            nodeNames,
            totalCost,
            size(nodeIdList) - 1 AS hops
        ORDER BY totalCost
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                source_id=source_id,
                target_id=target_id,
                k=k,
                weight_prop=weight_property
            )
            return [record.data() for record in result]
    
    # =========================================================
    # A* ALGORITHM
    # =========================================================
    
    def astar_path(
        self,
        source_id: str,
        target_id: str,
        latitude_property: str = "latitude",
        longitude_property: str = "longitude"
    ) -> Dict[str, Any]:
        """
        Find shortest path using A* algorithm.
        
        A* uses a heuristic (geographic distance) to guide search.
        Faster than Dijkstra when heuristic is available.
        
        Args:
            source_id: Starting node ID
            target_id: Target node ID
            latitude_property: Property with latitude
            longitude_property: Property with longitude
            
        Returns:
            Shortest path using A*
            
        Use Cases:
            - Geographically-aware routing
            - Finding nearest provider
            - Logistics optimization
        """
        query = """
        MATCH (source {id: $source_id}), (target {id: $target_id})
        CALL gds.shortestPath.astar.stream($graph_name, {
            sourceNode: source,
            targetNode: target,
            latitudeProperty: $lat_prop,
            longitudeProperty: $lon_prop
        })
        YIELD index, sourceNode, targetNode, totalCost, nodeIds, costs
        WITH totalCost,
             [nodeId IN nodeIds | gds.util.asNode(nodeId).id] AS nodeIdList,
             [nodeId IN nodeIds | gds.util.asNode(nodeId).name] AS nodeNames
        RETURN 
            nodeIdList,
            nodeNames,
            totalCost,
            size(nodeIdList) - 1 AS hops
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                source_id=source_id,
                target_id=target_id,
                lat_prop=latitude_property,
                lon_prop=longitude_property
            )
            record = result.single()
            return record.data() if record else None
    
    # =========================================================
    # BREADTH-FIRST SEARCH
    # =========================================================
    
    def bfs_paths(
        self,
        source_id: str,
        max_depth: int = 3
    ) -> List[Dict[str, Any]]:
        """
        Explore graph using Breadth-First Search.
        
        Finds all reachable nodes level by level.
        
        Args:
            source_id: Starting node ID
            max_depth: Maximum depth to explore
            
        Returns:
            Nodes at each depth level
        """
        query = """
        MATCH (source {id: $source_id})
        CALL gds.bfs.stream($graph_name, {
            sourceNode: source,
            maxDepth: $max_depth
        })
        YIELD path
        WITH path
        UNWIND range(0, size(nodes(path))-1) AS idx
        WITH nodes(path)[idx] AS node, idx AS depth
        RETURN 
            depth,
            collect(DISTINCT {id: node.id, name: node.name, labels: labels(node)}) AS nodesAtDepth,
            count(DISTINCT node) AS nodeCount
        ORDER BY depth
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                source_id=source_id,
                max_depth=max_depth
            )
            return [record.data() for record in result]
    
    # =========================================================
    # SINGLE SOURCE SHORTEST PATH
    # =========================================================
    
    def single_source_shortest_paths(
        self,
        source_id: str,
        weight_property: str = None
    ) -> List[Dict[str, Any]]:
        """
        Find shortest paths from one node to all others.
        
        Useful for understanding overall reachability.
        
        Args:
            source_id: Starting node ID
            weight_property: Optional weight property
            
        Returns:
            Shortest paths to all reachable nodes
        """
        query = """
        MATCH (source {id: $source_id})
        CALL gds.allShortestPaths.dijkstra.stream($graph_name, {
            sourceNode: source,
            relationshipWeightProperty: $weight_prop
        })
        YIELD index, sourceNode, targetNode, totalCost, nodeIds
        WITH totalCost,
             gds.util.asNode(targetNode).id AS targetId,
             gds.util.asNode(targetNode).name AS targetName,
             labels(gds.util.asNode(targetNode))[0] AS targetLabel,
             size(nodeIds) - 1 AS hops
        WHERE targetId <> $source_id
        RETURN targetId, targetName, targetLabel, totalCost, hops
        ORDER BY totalCost
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                graph_name=self.graph_name,
                source_id=source_id,
                weight_prop=weight_property
            )
            return [record.data() for record in result]
    
    def cleanup(self):
        """Drop the graph projection."""
        query = f"CALL gds.graph.drop('{self.graph_name}', false)"
        with self.driver.session() as session:
            try:
                session.run(query)
            except:
                pass


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Pathfinding Algorithms - Demonstration")
    print("=" * 60)
    
    print("""
    Pathfinding Algorithms in GDS:
    
    1. DIJKSTRA: Single shortest path with weights
       - Best for: Finding optimal route with costs
       - Complexity: O((V+E) log V)
    
    2. ALL SHORTEST PATHS: All paths of minimum length
       - Best for: Finding alternative equal routes
       - Complexity: O(V * E)
    
    3. YEN'S K-SHORTEST: Top K different paths
       - Best for: Finding backup/alternative routes
       - Complexity: O(K * V * (E + V log V))
    
    4. A*: Heuristic-guided shortest path
       - Best for: Geographic routing
       - Complexity: O(E) with good heuristic
    
    5. BFS: Level-by-level exploration
       - Best for: Finding all reachable nodes
       - Complexity: O(V + E)
    """)
    
    try:
        driver = get_driver()
        pathfinder = PathfindingAlgorithms(driver)
        
        # Create projection
        print("\n📌 Creating Graph Projection")
        print("-" * 40)
        proj = pathfinder.create_graph_projection(include_weights=False)
        if proj:
            print(f"   Nodes: {proj['nodeCount']}")
            print(f"   Relationships: {proj['relationshipCount']}")
        
        # Example queries would go here with real data
        print("\n📌 Pathfinding Ready")
        print("-" * 40)
        print("   Use dijkstra_shortest_path() to find optimal paths")
        print("   Use k_shortest_paths() for alternatives")
        print("   Use bfs_paths() for exploration")
        
        pathfinder.cleanup()
        driver.close()
        
    except Exception as e:
        print(f"   Error: {e}")
    
    print("\n" + "=" * 60)
    print("✅ Pathfinding demonstration complete!")
    print("=" * 60)
